# ver 1.4 2009-05-27
# debug
# 2009-05-24 tosion initialization problem
# 2009-05-26 seperate the lennard jones potential into single lines in order to calculate the 14 interaction
#            seperate the coulomb substract potential into single lines in order to calculate the 14 interaction


import math
import shutil

POW_2_sixth = math.pow(2, 1.0/6)
Kcal_2_eV = 23.061
ELEMENT_TABLE = {"H":1.008, "C":12.011, "O":15.9994, "N":14.01}

AMBER_vdw_14scale = 0.5
AMBER_charge_14scale = 0.833
cutoff_14_charge = 8.0

ATOM_TYPE_EQT = {"H1":"H[H2O]" , "O1":"OH2"}
chargeparameter = {"H1":0.4 , "O1":-0.8}

def TOSP2RB(TCOSPF, package = "gromacs"):
    f = TCOSPF
    c = [0, 0, 0, 0, 0, 0]
    c[0] = 2*f[2] + f[1] + f[3]
    c[1] = -f[1] + 3*f[3]
    c[2] = -2*f[2] + 8*f[4]
    c[3] = -4*f[3]
    c[4] = -8*f[4]
    c[5] = 0

    if package == "gulp":
        c = [c[i]*math.pow(-1,i) for i in range(len(c))]
    elif package == "gromacs":
        c = c

    return c

def combine_arithmetic(epsilon1, sigma1, epsilon2, sigma2):
    epsilon = math.sqrt(float(sigma1) * float(sigma2))
    sigma = 0.5 * (float(epsilon1) + float(epsilon2))
    return epsilon, sigma

def getbondparameters(ppffile):
    bondparameters = []
    
    f = open(ppffile, 'r')
    for i in f:
        if i.startswith("#"):
            pass
        else:
            term, atomtype, forcefieldparameters = i.split(":", 3)
            term = [term]
            atomtype = [i.strip() for i in atomtype.split(',')]
            forcefieldparameters = [i.strip("* \t\n\r") for i in forcefieldparameters.split(",")]                
            bondparameters.append(term + atomtype + forcefieldparameters)
            
    return bondparameters
            

def writeSpecies(o, inputff, eqt):
    specie_type = []
    o.write("species\n")
    for i in inputff:
        if i[0] == "ATYPE" or i[0] == "ATC":
            o.write("%-10score\n"%eqt[i[1]])
    o.write("\nelement\n")
    for i in inputff:
        if i[0] == "ATYPE"or i[0] == "ATC":
            specie_type.append(i[1])
            element = ELEMENT_TABLE[i[1][:1].upper()]
            o.write("mass %-10s%-10.5f\n"%(eqt[i[1]],element))
    o.write("end\n")

    return specie_type
                    

def writeBHARM(o, inputff, eqt, fit = 0):
    o.write("\nharmonic bond kcal\n")
    for i in inputff:
        if i[0] == "BHARM":
            o.write("%-10s%-10s%10.2f%10.4f"%(eqt[i[1]], eqt[i[2]], float(i[4])*2,float(i[3])))
            if fit == 1:
                o.write(" 0.0 0.0\n")
            else:
                o.write("\n")
    
def writeAHARM(o, inputff, eqt, fit= 0):
    o.write("\nthree-body bond kcal\n")
    for i in inputff:
        if i[0] == "AHARM":
            o.write("%-10s%-10s%-10s%10.2f%10.4f"%(eqt[i[2]], eqt[i[1]], eqt[i[3]], float(i[5])*2,float(i[4])))
            if fit == 1:
                o.write(" 0.0 0.0\n")
            else:
                o.write("\n")

def writeTCOSP(o, inputff, eqt, fit = 0):
    rcut = 8.0
    cutoff = "%4.1f%4.1f%4.1f%4.1f"%(rcut, rcut, rcut, rcut)
    o.write("\ntorsion ryck bond 3 single\n")
    for i in inputff:
        if i[0] == "TCOSP":
            o.write("%s"%"".join(["%-6s"%eqt[a] for a in i[1:5]]))
            TCOSPF = ["nonsense",0,0,0,0,0]
            
            for a in range(1,(((len(i)-4)/3)+1)):
                TCOSPF[int(float(i[3*a + 4]))] = float(i[3*a + 3])
            RB = TOSP2RB(TCOSPF, "gulp")
            RB = [ i/Kcal_2_eV for i in RB ]
            o.write("%10.4f"%RB[0])
            if fit == 1:
                o.write("%s 0.0 0.0 0.0 0.0\n"%cutoff)
            else:
                o.write("\n")
            o.write("%s\n"%"".join(["%-9.4f"%a for a in RB[1:4]]))
 
def writeN12_6_geometric(o, inputff, eqt):
    rmin = 0
    rmax = 12.0
    
    o.write("\nepsilon kcal\n")
    for i in inputff:
        if i[0] == "N12_6":
            o.write("%-10s%10.4f%10.4f\n"%(i[1], float(i[3]),float(i[2])))
    o.write("\nlennard epsilon 12 6 inter kcal geometric all\n")
    o.write("%-10.2f%-10.2f\n"%(rmin, rmax))
    o.write("\nlennard epsilon 12 6 intra kcal o14 geometric all %4.2f\n"%AMBER_vdw_14scale)
    o.write("%-10.2f%-10.2f\n"%(rmin, rmax))


def writeN12_6_explicit(o, inputff, eqt, fit = 0):
    vdw_parameters = []
    rmin = 0
    rmax = 8.0
    
    for i in inputff:
        if i[0] == "N12_6":
            vdw_parameters.append(i)
    for i in range(len(vdw_parameters)):
        for j in range( i, len(vdw_parameters)):
            epsilon, sigma = combine_arithmetic(vdw_parameters[i][2],vdw_parameters[i][3],vdw_parameters[j][2],vdw_parameters[j][3])
            o.write("\nlennard epsilon 12 6 kcal x13 %4.2f\n"%AMBER_vdw_14scale)
            o.write("%-10s%-10s%10.4f%10.4f"%(eqt[vdw_parameters[i][1]],eqt[vdw_parameters[j][1]], epsilon, sigma))
            o.write("%6.1f%6.1f"%(rmin, rmax))
            if fit == 1:
                o.write(" 0.0 0.0 \n")
            else:
                o.write("\n")

            
def write_charge_1_4(o, specie_type, eqt, fit = 0):

    for i in range(len(specie_type)):
        for j in range(i,len(specie_type)):
            o.write("\ncoulomb intra o14 %f\n"%(1.0 - AMBER_charge_14scale))
            o.write("%-10s%-10s%6.2f"%(eqt[specie_type[i]], eqt[specie_type[j]],cutoff_14_charge))
            if fit == 1:
                o.write(" 0.0\n")
            else:
                o.write("\n")
    

def ppf2lib(dffppf, msdfile, eqt_table_fftype_apptype, fit = 0):

    fit = int(fit)
    
    if fit == 0:
        o = open("%s.lib"%dffppf, 'w')
        parameters = getbondparameters("%s.ppf"%dffppf)
    
        specie_type = writeSpecies(o, parameters,eqt_table_fftype_apptype )
        writeBHARM(o, parameters, eqt_table_fftype_apptype)
        writeAHARM(o, parameters, eqt_table_fftype_apptype)
        writeTCOSP(o, parameters, eqt_table_fftype_apptype)
        writeN12_6_explicit(o, parameters, eqt_table_fftype_apptype)
        write_charge_1_4(o, specie_type, eqt_table_fftype_apptype)
    
        o.close()
    else:
        o = open("%s.gin"%msdfile, 'a')
        parameters = getbondparameters("%s.ppf"%dffppf)
    
        specie_type = writeSpecies(o, parameters,eqt_table_fftype_apptype)
        writeBHARM(o, parameters, eqt_table_fftype_apptype, 1)
        writeAHARM(o, parameters, eqt_table_fftype_apptype, 1)
        writeTCOSP(o, parameters, eqt_table_fftype_apptype, 1)
        writeN12_6_explicit(o, parameters, eqt_table_fftype_apptype, 1)
        write_charge_1_4(o, specie_type, eqt_table_fftype_apptype, 1)

        o.close()
    
if __name__ == "__main__":
    
    ppf2lib("water")
                    
